package com.liezh.client;

import com.liezh.SessionUtil;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.*;
import org.apache.thrift.transport.TTransport;

import java.io.PrintStream;
import java.util.Map;

/**
 * @author liezh
 * @date 2022-3-7
 */
public class THeaderClientProtocol extends TBinaryProtocol {

    public THeaderClientProtocol(TTransport transport) {
        super(transport);
    }

    //clientSend
    @Override
    public void writeMessageBegin(TMessage tMessage) throws TException {
        String methodName = tMessage.name;
        TTransport transport = this.getTransport();

        super.writeMessageBegin(tMessage);
        //write header to field0
        writeFieldZero();
        final PrintStream format = System.out.format("methodName %s add header %s", methodName, SessionUtil.getSession().toString());
        format.println();
    }


    public void writeFieldZero() throws TException {
        TField TRACE_HEAD = new TField("THeader", TType.MAP, (short) 0);
        this.writeFieldBegin(TRACE_HEAD);
        {
            Map<String, String> headerInfo = genHeaderInfo();
            this.writeMapBegin(new TMap(TType.STRING, TType.STRING, headerInfo.size()));
            for (Map.Entry<String, String> entry : headerInfo.entrySet()) {
                this.writeString(entry.getKey());
                this.writeString(entry.getValue());
            }
            this.writeMapEnd();
        }
        this.writeFieldEnd();
    }

    private Map<String, String> genHeaderInfo() {
        //gen header infoSessionUtil.getSession();
        return SessionUtil.getSession();
    }

    //clientReceive
    @Override
    public TMessage readMessageBegin() throws TException {
        TMessage tMessage = super.readMessageBegin();
        System.out.println("type " + tMessage.type + " name " + tMessage.name);
        return tMessage;
    }
}
